import asyncio
import nest_asyncio
import sys
import os
import numpy as np  # 新增：导入numpy库
from lightrag import LightRAG, QueryParam
from lightrag.llm.openai import openai_complete_if_cache, openai_embed
from lightrag.utils import EmbeddingFunc
from dotenv import load_dotenv

load_dotenv()

WORKING_DIR = "./dickens"

if not os.path.exists(WORKING_DIR):
    os.mkdir(WORKING_DIR)


async def llm_model_func(prompt, system_prompt=None, history_messages=[], keyword_extraction=False, **kwargs):
    return await openai_complete_if_cache(
        os.getenv("LLM_MODEL"),
        prompt,
        system_prompt=system_prompt,
        history_messages=history_messages,
        api_key=os.getenv("LLM_BINDING_API_KEY"),
        base_url=os.getenv("LLM_BINDING_HOST"),
        **kwargs,
    )


async def embedding_func(texts: list[str]) -> np.ndarray:
    return await openai_embed(
        texts,
        model=os.getenv("EMBEDDING_MODEL"),
        api_key=os.getenv("EMBEDDING_BINDING_API_KEY"),
        base_url=os.getenv("EMBEDDING_BINDING_HOST"),
    )


async def get_embedding_dim():
    test_text = ["This is a test sentence."]
    embedding = await embedding_func(test_text)
    return embedding.shape[1]


async def initialize_rag():
    embedding_dimension = await get_embedding_dim()
    print(f"Detected embedding dimension: {embedding_dimension}")

    rag = LightRAG(
        working_dir=WORKING_DIR,
        graph_storage=os.getenv("LIGHTRAG_GRAPH_STORAGE"),
        kv_storage=os.getenv("LIGHTRAG_KV_STORAGE"),
        vector_storage=os.getenv("LIGHTRAG_VECTOR_STORAGE"),
        doc_status_storage=os.getenv("LIGHTRAG_DOC_STATUS_STORAGE"),
        llm_model_func=llm_model_func,
        embedding_func=EmbeddingFunc(
            embedding_dim=embedding_dimension,
            max_token_size=8192,
            func=embedding_func,
        ),
    )

    await rag.initialize_storages()
    return rag


async def main():
    # try:
    rag = await initialize_rag()
    #
    # custom_kg = {
    #     "chunks": [
    #         {
    #             "content": "Alice和Bob正在合作进行量子计算研究。",
    #             "source_id": "doc-1"
    #         }
    #     ],
    #     "entities": [
    #         {
    #             "entity_name": "Alice",
    #             "entity_type": "person",
    #             "description": "Alice是一位专门研究量子物理的研究员。",
    #             "source_id": "doc-1"
    #         },
    #         {
    #             "entity_name": "Bob",
    #             "entity_type": "person",
    #             "description": "Bob是一位数学家。",
    #             "source_id": "doc-1"
    #         },
    #         {
    #             "entity_name": "量子计算",
    #             "entity_type": "technology",
    #             "description": "量子计算利用量子力学现象进行计算。",
    #             "source_id": "doc-1"
    #         }
    #     ],
    #     "relationships": [
    #         {
    #             "src_id": "Alice",
    #             "tgt_id": "Bob",
    #             "description": "Alice和Bob是研究伙伴。",
    #             "keywords": "合作 研究",
    #             "weight": 1.0,
    #             "source_id": "doc-1"
    #         },
    #         {
    #             "src_id": "Alice",
    #             "tgt_id": "量子计算",
    #             "description": "Alice进行量子计算研究。",
    #             "keywords": "研究 专业",
    #             "weight": 1.0,
    #             "source_id": "doc-1"
    #         },
    #         {
    #             "src_id": "Bob",
    #             "tgt_id": "量子计算",
    #             "description": "Bob研究量子计算。",
    #             "keywords": "研究 应用",
    #             "weight": 1.0,
    #             "source_id": "doc-1"
    #         }
    #     ]
    # }
    #
    # rag.insert_custom_kg(custom_kg)  # 添加await关键字
    #
    # # 清除所有缓存
    #
    # # await rag.delete_by_entity("Google")
    # entity = await rag.acreate_entity("Google", {
    #     "description": "Google是一家专注于互联网相关服务和产品的跨国科技公司。",
    #     "entity_type": "company"
    # })
    # # 创建另一个实体
    # product = await rag.acreate_entity("Gmail", {
    #     "description": "Gmail是由Google开发的电子邮件服务。",
    #     "entity_type": "product"
    # })
    # # print(product)
    # await rag.aclear_cache()
    # 创建实体之间的关系
    # relation = rag.create_relation("Google", "Gmail", {
    #     "description": "Google开发和运营Gmail。",
    #     "keywords": "开发 运营 服务",
    #     "weight": 2.0
    # })

    # 编辑现有实体
    updated_entity = rag.edit_entity("路由器组", {
        # "description": "Google是Alphabet Inc.的子公司，成立于1998年。",
        # "entity_type": "tech_company",
        "file_path": "测试修改文件.txt",

    })

    await rag.aclear_cache()




# except Exception as err:
#     print(f"An error occurred: {err}")
# finally:
#     if hasattr(rag, 'finalize_storages'):
#         await rag.finalize_storages()

if __name__ == "__main__":
    nest_asyncio.apply()
    loop = asyncio.get_event_loop()
    loop.run_until_complete(main())
